package com.feidee.fd.sml.algorithm.forecast

import org.apache.spark.ml.{PipelineModel, PipelineStage}

/**
  * @Author songhaicheng
  * @Date 2019/4/23 11:09
  * @Description 通过类型匹配和下标，自动查找模型 PipelineStages 中的目标
  * @Reviewer
  */
class StageFinder[T: Manifest](model: PipelineModel) {

  private[algorithm] def findWithOrder(order: Int = 0): Option[T] = {
    var idx = 0
    for (stage: PipelineStage <- model.stages) {
      stage match {
        case s: T =>
          if (order == idx) {
            return Some(s.asInstanceOf[T])
          }
          else {
            idx += 1
          }
        case _ =>
      }
    }
    None
  }
}
class PSStageFinder[T: Manifest](model: com.tencent.angel.sona.ml.PipelineModel) {

  private[algorithm] def findWithOrder(order: Int = 0): Option[T] = {
    var idx = 0
    for (stage: com.tencent.angel.sona.ml.PipelineStage <- model.stages) {
      stage match {
        case s: T =>
          if (order == idx) {
            return Some(s.asInstanceOf[T])
          }
          else {
            idx += 1
          }
        case _ =>
      }
    }
    None
  }
}
